Source: LucidJS/src/optvis/objectives.js

import * as tf from '@tensorflow/tfjs';

/**
 * Visualize a single channel.
 */
export function channel(model, options) {
 const layerOutput = model.getLayer(options.layer).output;
 const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
 const mul = options.neg ? 1 : -1;
 const s = layerOutput.shape;
 const ch = Math.min(options.channel, s[3]-1);

  const inner =  (inputs) => {
    return tf.mean(
        auxModel.apply(inputs, {training: true})
        .gather([ch], 3).mul(mul)
    );
  }
  return inner;
}

/**
 * Maximize 'interestingness' at some layer.
 *
 *
 * See Mordvintsev et al., 2015.
 */
export function deepdream(model, options) {
   const layerOutput = model.getLayer(options.layer).output;
   const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
   const mul = options.neg ? 1 : -1;

    const inner =  (inputs) => {
      return tf.mean(
          auxModel.apply(inputs, {training: true}).pow(2).mul(mul)
      );
    }
    return inner;
}

/**
 * Maximize output / class activation.
 * @param {*} model 
 * @param {*} options 
 */
export function output(model, options) {
   const layerOutput = model.outputs[0].inputs[0];
   const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
   const mul = options.neg ? 1 : -1;

    const inner =  (inputs) => {
      return tf.mean(
          auxModel.apply(inputs, {training: true})
          .reshape([1, -1]).slice([0, options.out], [1, 1]).mul(mul)
      );
    }
    return inner;
}

/**
* Visualize a single neuron of a single channel.
*
* Defaults to the center neuron. When width and height are even numbers, we
* choose the neuron in the bottom right of the center 2x2 neurons.
*
* Odd width & height:               Even width & height:
*
* +---+---+---+                     +---+---+---+---+
* |   |   |   |                     |   |   |   |   |
* +---+---+---+                     +---+---+---+---+
* |   | X |   |                     |   |   |   |   |
* +---+---+---+                     +---+---+---+---+
* |   |   |   |                     |   |   | X |   |
* +---+---+---+                     +---+---+---+---+
*                                   |   |   |   |   |
*                                   +---+---+---+---+
*/
export function neuron(model, options) {
   const layerOutput = model.getLayer(options.layer).output;
   const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
   const mul = options.neg ? 1 : -1;
   let x, y;
   const s = layerOutput.shape;
   if(!options.neuron) {
     x = Math.floor(s[1]/2);
     y = Math.floor(s[2]/2);
   } else {
     const [xn, yn] = options.neuron;
     x = Math.min(xn, s[2]-1);
     y = Math.min(yn, s[1]-1);
   }

   const ch = Math.min(options.channel, s[3]-1);

    const inner =  (inputs) => {
      return tf.mean(
        auxModel.apply(inputs, {training: true})
          .slice([0, y, x, ch], [1, 1, 1, 1])
          .mul(mul)
        );
    }
    return inner;
}

/**
 * Maximize single "pixel" location
 * @param {*} model 
 * @param {*} options 
 */
export function spatial(model, options) {

   const layerOutput = model.getLayer(options.layer).output;
   const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
   const mul = options.neg ? 1 : -1;
   let x, y;
   const s = layerOutput.shape;
   if(!options.neuron) {
     x = Math.floor(s[1]/2);
     y = Math.floor(s[2]/2);
   } else {
     const [xn, yn] = options.neuron;
     x = Math.min(xn, s[2]-1);
     y = Math.min(yn, s[1]-1);
   }

    const inner =  (inputs) => {
      return tf.mean(
        auxModel.apply(inputs, {training: true})
          .slice([0, y, x, 0], [1, 1, 1, s[3]-1])
          .mul(mul)
        );
    }
    return inner;
}

/**
 * Experimental style objective
 * @param {*} model 
 * @param {*} contentImg 
 * @param {*} styleImg 
 * @param {*} contentLrs 
 * @param {*} styleLrs 
 */
export function style(model, contentImg, styleImg, contentLrs, styleLrs) {

  const targetStyleActivations = getActivationsForLayers(
    model, styleImg, styleLrs);

  const targetContentActivations = getActivationsForLayers(
    model, contentImg, contentLrs);

  const styleAuxModel = getAuxModel(
    model, getLayerOutputs(model, styleLrs));

  const contentAuxModel = getAuxModel(
    model, getLayerOutputs(model, contentLrs));

  const inner = (inputs) => {
    const inputStyleActivations = makeArray(styleAuxModel.apply(inputs));
    const inputContentActivations = makeArray(contentAuxModel.apply(inputs));

    const contentObj = activationDifference(
      inputContentActivations, targetContentActivations, gramStyle).mul(100);

    const styleObj = activationDifference(
      inputStyleActivations, targetStyleActivations, gramStyle);

    console.log("content loss:", contentObj.dataSync(),
  "style loss:", styleObj.dataSync())


    const objective = contentObj.add(styleObj);
    return objective;
  }

  return inner;
}

/**
 * Difference between activations of two input images.
 * @param {*} imgActivations 
 * @param {*} targetActivations 
 * @param {*} transformF 
 * @param {*} activLossF 
 */
function activationDifference(
  imgActivations, targetActivations, transformF=null, activLossF=meanL1Loss){
  const losses = [];
  for(let i=0; i< imgActivations.length; i++) {
    let iA = imgActivations[i];
    let tA = targetActivations[i];
    if(transformF) {
      iA = transformF(iA);
      tA = transformF(tA);
    }
    const loss = activLossF(iA, tA);
    losses.push(loss);
  }
  return tf.addN(losses);
}

/**
 * Returns list of activation tensors for each layer.
 * @param {*} model 
 * @param {*} image 
 * @param {*} layers 
 */
function getActivationsForLayers(model, image, layers) {
  const outputs = getLayerOutputs(model, layers);

  const auxModel = getAuxModel(model, outputs);
  let activations = auxModel.apply(image);

  activations = makeArray(activations);

  const activationList = [];
  activations.forEach((tensor) => {
    const shp = tensor.shape;
    const data = tensor.dataSync();
    tensor.dispose();
    const frozenTensor = tf.tensor(data, shp);
    activationList.push(frozenTensor);
  });

  return activationList;
}

/**
 * Returns new model with specified outputs.
 * @param {*} model 
 * @param {*} outputs 
 */
function getAuxModel(model, outputs){
  const auxModel = tf.model(
    {inputs: model.inputs, outputs: outputs});
  return auxModel;
}

/**
 * Returns output layer tensors.
 * @param {*} model 
 * @param {*} layers 
 */
function getLayerOutputs(model, layers) {
  const outputs = [];
  layers.forEach((layerName) => {
    const outLayer = model.getLayer(layerName);
    outputs.push(outLayer.output);
  });
  return outputs;
}

/**
 * Returns array of tensors from tensor or array of tensors
 * @param {*} tensorOrArray 
 */
function makeArray(tensorOrArray) {
  if(!Array.isArray(tensorOrArray)){
    tensorOrArray  = [tensorOrArray];
  }
  return tensorOrArray;
}

/**
 * Optimize for specified activation modifications.
 * @param {*} model 
 * @param {*} originalImage 
 * @param {*} activationModDict 
 * @param {*} activationLossF 
 */
export function activationModification(model, originalImage,
  activationModDict, activationLossF=meanL2Loss){
  const layerOutputs = [];
  const layerNames = [];
  for(const [layerName, val] of Object.entries(activationModDict)){
    const outLayer = model.getLayer(layerName);
    layerOutputs.push(outLayer.output);
    layerNames.push(layerName);
  }
  if(!layerOutputs.length){
    return (inputs) => tf.mean(inputs.mul(0));
  }

  const auxModel = tf.model(
    {inputs: model.inputs, outputs: layerOutputs});

  let origImageActivations = auxModel.apply(originalImage);
  const origActivationList = getActivationsForLayers(model, originalImage, layerNames);

  const inner = (inputs) => {

    // we also set get the activations of the optimized image which will change during optimization
    let currentActivations = auxModel.apply(inputs);

    // we use the supplied loss function to compute the actual losses
    let losses = [];

    currentActivations = makeArray(currentActivations);

    for(let i=0; i<origActivationList.length; i++){
      const oA = origActivationList[i];
      let cA = currentActivations[i];
      const [b, h, w, ch] = cA.shape;
      const layerName = layerNames[i];
      const channelList = [];
      let shiftList = [];
      let scaleList = [];
      const noiseList = [];
      for(let j=0; j<ch; j++) {
        if(j in activationModDict[layerName]){
          const values = activationModDict[layerName][j];
          channelList.push(channel);
          shiftList.push(values['shift']);
          scaleList.push(values['scale']);
          noiseList.push(values['noise']);
        } else {
          shiftList.push(0);
          scaleList.push(1);
          noiseList.push(0);
        }
      }

      scaleList = [[[scaleList]]];
      const scaleT = tf.tensor(scaleList);
      cA = cA.div(scaleT);

      shiftList = [[[shiftList]]];
      const shiftT = tf.tensor(shiftList);
      cA = cA.sub(shiftT);

      const layerLoss = activationLossF(oA, cA);
      losses.push(layerLoss);
    }

    const loss = tf.addN(losses);
    console.log("act diff loss: ", loss.dataSync())
    return loss;
  }

  return inner;
}

/**
 * Returns gram matrix of input tensor, normalized by length of flat length.
 * @param {*} array 
 */
function gramStyle(array){
  const sh = array.shape;
  const channels = sh[sh.length-1];
  const array_flat = tf.reshape(array, [-1, channels]);
  const length = array_flat.shape[0];
  const gram = tf.matMul(array_flat, array_flat, true, false);
  return gram.div(length);
}

/**
 * Returns mean of absolute differences.
 * @param {*} g1 
 * @param {*} g2 
 */
function meanL1Loss(g1, g2) {
  return tf.mean(tf.abs(g1.sub(g2)));
}

/**
 * Returns mean of squared differences.
 * @param {*} g1 
 * @param {*} g2 
 */
function meanL2Loss(g1, g2) {
  return tf.mean(tf.pow(g1.sub(g2), 2));
}